import numpy as np

def pso(x, y, option, data):
    """
    粒子群优化算法实现
    参数:
        x: 初始粒子位置矩阵 (numAgent × dim)
        y: 初始适应度值向量 (numAgent × 1)
        option: 算法参数配置字典
        data: 问题数据字典
    返回:
        bestY: 最优适应度值
        bestX: 最优粒子位置
        recording: 优化过程记录字典
    """
    # 初始化记录字典
    recording = {
        'bestFit': np.zeros(option['maxIteration'] + 1),
        'meanFit': np.zeros(option['maxIteration'] + 1)
    }
    
    # 初始化速度
    v = np.random.randn(*x.shape)
    
    # 初始化全局最优和个体最优
    y_g = np.min(y)
    position = np.argmin(y)
    x_g = x[position, :].copy()
    y_p = y.copy()
    x_p = x.copy()
    
    # 记录初始最优值
    recording['bestFit'][0] = y_g
    recording['meanFit'][0] = np.mean(y_p)
    
    # 提取算法参数
    w_pso = option['w_pso']

    lb = option['lb']
    ub = option['ub']
    fobj = option['fobj']
    numAgent = option['numAgent']
    dim = option['dim']
    
    # 初始化停滞检测变量
    stagnation_count = 0
    prev_best = y_g
    stagnation_threshold = option.get('stagnation_threshold', 20)  # 停滞阈值，可通过option配置
    
    # 初始化早停参数
    early_stop_count = 0
    epsilon = option.get('epsilon', 1e-6)  # 最小改进阈值
    early_stop_threshold = option.get('early_stop_threshold', 15)  # 早停阈值
    
    # 初始化打印参数
    print_interval = option.get('print_interval', 10)

    # 主迭代过程
    for iter in range(1, option['maxIteration'] + 1):
        if iter % print_interval == 0 or iter == 1:
            print(f"PSO, iter: {iter}, minFit: {y_g:.6f}, stagnation_count: {stagnation_count}")
        
        # 生成随机数
        r1 = np.random.rand(numAgent, dim)
        r2 = np.random.rand(numAgent, dim)

        # 改进的自适应参数策略
        # 1. 惯性权重：非线性递减策略
        current_w = 0.4 + (0.9 - 0.4) * np.exp(-3 * (iter / option['maxIteration'])**2)

        # 2. 时变加速系数：c1从2.5递减到0.5，c2从0.5递增到2.5
        c1 = 2.5 - 2 * (iter / option['maxIteration'])
        c2 = 0.5 + 2 * (iter / option['maxIteration'])

        # 3. 速度钳制系数
        max_vel = 0.3 * (ub - lb)

        # 更新所有粒子速度（向量化操作）
        v = current_w * v + c1 * r1 * (x_p - x) + c2 * r2 * (x_g - x)
        v = np.clip(v, -max_vel, max_vel)
        
        # 更新所有粒子位置
        x += v
        
        # 边界检查（向量化）
        x = np.clip(x, lb, ub)
        
        # 计算所有粒子适应度
        y = np.array([fobj(x[i, :], option, data)['fit'] for i in range(numAgent)])
        
        # 更新个体最优
        improved = y < y_p
        y_p[improved] = y[improved]
        x_p[improved] = x[improved].copy()
        
        # 更新全局最优
        current_min_idx = np.argmin(y_p)
        current_min = y_p[current_min_idx]
        if current_min < y_g:
            y_g = current_min
            x_g = x_p[current_min_idx].copy()
        
        # 停滞检测与处理
        if y_g < prev_best:
            prev_best = y_g
            stagnation_count = 0
        else:
            stagnation_count += 1
            if stagnation_count >= stagnation_threshold:
                # 重置速度以增加多样性
                v = 0.8 * np.random.randn(numAgent, dim) * max_vel
                stagnation_count = 0  # 重置停滞计数
        
        # 早停检查
        current_improvement = prev_best - y_g
        if current_improvement < epsilon:
            early_stop_count += 1
            if early_stop_count >= early_stop_threshold:
                print(f"Early stopping at iteration {iter} as improvement ({current_improvement:.8f}) is below threshold {epsilon} for {early_stop_threshold} iterations.")
                break
        else:
            early_stop_count = 0
        
        # 记录优化过程
        recording['bestFit'][iter] = y_g
        recording['meanFit'][iter] = np.mean(y_p)
    
    # 截断记录数组至实际迭代次数
    recording['bestFit'] = recording['bestFit'][:iter+1]
    recording['meanFit'] = recording['meanFit'][:iter+1]

    return y_g, x_g, recording